//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////
// gTraySpy.cpp
//
// Main implementation file for gTraySpy.dll
//
// An slimmed-down, optimized version of geOTraySpy.dll for balloon
// notification specific hooking.
//
// geOTraySpy is based on Mike Lin's traysaver code.
//
// It builds a dll which contains code to establish a hook into
// the Shell_TrayWnd class, receiving all wndProc messages after they
// have been processed by Shell_TrayWnd itself.
//
// It then forwards any copydata messages to SnarlTray.
//
//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////

#include <windows.h>
#include <stdio.h>
#include <TCHAR.H>
#include "gTraySpy.h"

//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////
// Globals
//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////

#define PAD_SIZE    65

struct stTrayItem
{
    COPYDATASTRUCT  CopyData;

    struct stNicHolder
    {
        BYTE            bPad[8];
        NOTIFYICONDATA  IconData;
        CHAR            szPad[PAD_SIZE];    // The Tip can be up to 128 chars long
    }               NicHolder;
};

#define MAX_SAVEDTRAYDATA   25

#pragma data_seg( "shared" )

CHAR                g_cServerClass [100]    = "";
HWND                g_hTargetWin            = 0;
HWND                g_hTrayWin              = 0;
HHOOK               g_hHook                 = 0;
UINT                g_iNextFreePos          = 0;
struct stTrayItem   g_SavedTrayData[MAX_SAVEDTRAYDATA];

#pragma data_seg()

#pragma comment( linker, "/SECTION:shared,rws" )

//const UINT WM_GTRAYSPY_MSG    = ::RegisterWindowMessage(GTRAYSPY_MSG);
const char CLASS_TrayProc[]     = "Shell_TrayWnd";

#ifndef NIF_INFO
#define NIF_INFO         0x00000010
#endif

typedef struct _MYNOTIFYICONDATA {
    DWORD   cbSize; 
    HWND    hWnd; 
    UINT    uID; 
    UINT    uFlags; 
    UINT    uCallbackMessage; 
    HICON   hIcon; 
    TCHAR   szTip[128];
    DWORD   dwState;
    DWORD   dwStateMask;
    TCHAR   szInfo[256];
    union
    {
        UINT  uTimeout;
        UINT  uVersion;
    } MYDUMMYUNIONNAME;
    TCHAR   szInfoTitle[64];
    DWORD   dwInfoFlags;
} MYNOTIFYICONDATA, *PMYNOTIFYICONDATA;

//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////
// Functions
//////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////
// IsNT
//
// Returns true if we are running under Windows NT.
//////////////////////////////////////////////////////////////////////

BOOL IsNT()
{
    static BOOL         bNT = FALSE;
    static BOOL         bFound = FALSE;
    OSVERSIONINFO       ovi;
    
    if( !bFound )
    {
        ovi.dwOSVersionInfoSize = sizeof( OSVERSIONINFO );

        GetVersionEx( &ovi );

        bNT = ovi.dwPlatformId == VER_PLATFORM_WIN32_NT;

        bFound = TRUE;
    }

    return bNT;
}

//////////////////////////////////////////////////////////////////////
// TrayWndProc
//
// New window proc for the system tray window
//////////////////////////////////////////////////////////////////////

LRESULT CALLBACK TrayWndProc( int nCode, WPARAM wParam, LPARAM lParam )
{
    CWPRETSTRUCT *pInfo = (CWPRETSTRUCT *)( lParam );
    
    if( nCode >= 0 )
    {
        // Interested in copydata messages
        if ( pInfo->message == WM_COPYDATA )
        {            
			// Make sure we have a valid handle on the Tray window
            if( !g_hTrayWin || !IsWindow( g_hTrayWin ) )
                g_hTrayWin = FindWindow( CLASS_TrayProc, NULL );

            // Make sure this message is from the Tray window
            if (IsWindow( g_hTrayWin) && pInfo->hwnd == g_hTrayWin )
            {
                COPYDATASTRUCT* pcds = (COPYDATASTRUCT*)( pInfo->lParam );
				
                if( pcds->dwData == 1 ) // if this data refers to a tray icon...
                {
                    // Maintain our own copy of the tray icons
                    NOTIFYICONDATA  *nicData = (NOTIFYICONDATA *)(((BYTE *)pcds->lpData) + 8);
                    INT             iTrayCmd = *(INT *)(((BYTE *)pcds->lpData) + 4);
					UINT    i = 0;
                    if ((int)nicData->uID == -1)
                    {
                        return CallNextHookEx( g_hHook, nCode, wParam, lParam );
                    }

                    if (iTrayCmd == NIM_ADD || iTrayCmd == NIM_MODIFY)
                    {
                        BOOL    bFound = FALSE;
                        

                        for (i = 0; i < g_iNextFreePos && i < MAX_SAVEDTRAYDATA; i++)
                        {
                            if (g_SavedTrayData[i].NicHolder.IconData.uID == nicData->uID && 
                                g_SavedTrayData[i].NicHolder.IconData.hWnd == nicData->hWnd)
                            {
                                memcpy (&g_SavedTrayData[i].CopyData, pcds, sizeof (COPYDATASTRUCT));
                                if (iTrayCmd == NIM_ADD)
                                {
                                    memcpy (&g_SavedTrayData[i].NicHolder.IconData,
                                            nicData,
                                            (int) nicData->cbSize <= (sizeof (NOTIFYICONDATA) + PAD_SIZE) ? (int) nicData->cbSize : (sizeof (NOTIFYICONDATA) + PAD_SIZE));
                                }
                                else
                                {
                                    if ((nicData->uFlags | NIF_ICON && nicData->hIcon) || 
                                        (g_SavedTrayData[i].NicHolder.IconData.hIcon == NULL && nicData->hIcon))
                                    {
//                                      MessageBox(NULL, "Icon", "", MB_OK);
                                        g_SavedTrayData[i].NicHolder.IconData.hIcon = nicData->hIcon;
                                    }

                                    if ((nicData->uFlags | NIF_MESSAGE) ||
                                        (!g_SavedTrayData[i].NicHolder.IconData.uCallbackMessage && nicData->uCallbackMessage))
                                    {
//                                      MessageBox(NULL, "Msg", "", MB_OK);
                                        g_SavedTrayData[i].NicHolder.IconData.uCallbackMessage = nicData->uCallbackMessage;
                                    }

                                    if ((nicData->uFlags | NIF_TIP) ||
                                        (g_SavedTrayData[i].NicHolder.IconData.szTip[0] == '\0' && nicData->szTip[0] != '\0'))
                                    {
                                        memcpy (g_SavedTrayData[i].NicHolder.IconData.szTip, nicData->szTip, sizeof (nicData->szTip) + (int) nicData->cbSize - sizeof(NOTIFYICONDATA));
                                    }   

                                    g_SavedTrayData[i].NicHolder.IconData.uFlags = nicData->uFlags;
                                }

                                bFound = TRUE;
                                break;
                            }
                        }

                        if (!bFound && (i < MAX_SAVEDTRAYDATA))
                        {
                            memcpy (&g_SavedTrayData[i].CopyData, pcds, sizeof (COPYDATASTRUCT));
                            memcpy (&g_SavedTrayData[i].NicHolder.IconData,
                                    nicData,
                                    (int) nicData->cbSize <= (sizeof (NOTIFYICONDATA) + PAD_SIZE) ? (int) nicData->cbSize : (sizeof (NOTIFYICONDATA) + PAD_SIZE));
                            g_iNextFreePos++;
                        }
                    }
                    else if (iTrayCmd == NIM_DELETE)
                    {
                        for (i = 0; i < g_iNextFreePos && i < MAX_SAVEDTRAYDATA; i++)
                        {
                            if (g_SavedTrayData[i].NicHolder.IconData.uID == nicData->uID && 
                                g_SavedTrayData[i].NicHolder.IconData.hWnd == nicData->hWnd)
                            {
                                memmove(&g_SavedTrayData[i], &g_SavedTrayData[i + 1], sizeof(struct stTrayItem) * (g_iNextFreePos - (i + 1)));
                                g_iNextFreePos--;
                                break;
                            }
                        }
                    }
					
                    // Make sure we have a valid handle on the target window
                    if( !g_hTargetWin || !IsWindow( g_hTargetWin ) )
                        g_hTargetWin = FindWindow( g_cServerClass, NULL );

                    // OK?
                    if (IsWindow( g_hTargetWin ))
                    {
                        // Forward the message
                        SendMessage( g_hTargetWin, 
                                     WM_COPYDATA, 
                                     (WPARAM)( g_hTrayWin ), 
                                     pInfo->lParam );
						
                        //DWORD dw = GetLastError();
                        //if (dw>0){
							//char szMyString[14] ={0};
							//sprintf(szMyString, "0x%08x",dw);
							//MessageBox(NULL, szMyString, TEXT("Message Sent, LastError:"), MB_OK) ; }
                    }
                }
            }
        }
		/*
        else if ( pInfo->message == WM_GTRAYSPY_MSG )
        {
            if( !g_hTargetWin || !IsWindow( g_hTargetWin ) )
                g_hTargetWin = FindWindow( g_cServerClass, NULL );

            for (UINT i = 0; i < g_iNextFreePos && i < MAX_SAVEDTRAYDATA; i++)
            {
                g_SavedTrayData[i].CopyData.lpData = &g_SavedTrayData[i].NicHolder;

                SendMessage( g_hTargetWin, 
                             WM_COPYDATA, 
                             reinterpret_cast< WPARAM >( g_hTrayWin ), 
                             (long) &g_SavedTrayData[i] );
            }
        }
		*/
    }

    return CallNextHookEx( g_hHook, nCode, wParam, lParam );
}

//BOOL APIENTRY DllMain( HANDLE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved )
//{
//    return TRUE;
//}

//////////////////////////////////////////////////////////////////////
// SetTrayHook
//
// Sets our hook
//////////////////////////////////////////////////////////////////////


#if _WIN64
BOOL  SetTrayHook64 (HINSTANCE hInst)
#else
BOOL  SetTrayHook (HINSTANCE hInst)
#endif
{
    DWORD dwThreadID = 0;
    if( g_hHook )
        return FALSE;
//    if (strlen(szServerName) > 99 || strlen(szServerName) == 0)
//        return false;
    _tcscpy_s(g_cServerClass,_countof(_T("BalloonCapture")), _T("BalloonCapture"));
    if( IsNT() )
    {
        HWND hwndTray = FindWindow( CLASS_TrayProc, NULL );
        if( !hwndTray )
            return FALSE;
        if( hwndTray )
            dwThreadID = GetWindowThreadProcessId( hwndTray, NULL );
    }

    memset (g_SavedTrayData, '\0', MAX_SAVEDTRAYDATA * sizeof (COPYDATASTRUCT));
	#if _WIN64
	g_hHook = SetWindowsHookEx( WH_CALLWNDPROCRET, TrayWndProc, hInst, dwThreadID );
	#else
    g_hHook = SetWindowsHookEx( WH_CALLWNDPROCRET, TrayWndProc, hInst, dwThreadID );
	#endif
    //if (g_hHook == 0)
    //{
    //    DWORD dw = GetLastError(); 
    //    char szMyString[14] ={0};
    //    sprintf(szMyString, "0x%08x",dw);
    //}

    
    return g_hHook != 0;
}

//////////////////////////////////////////////////////////////////////
// UnsetTrayHook
//
// Unsets the hook
//////////////////////////////////////////////////////////////////////

#if _WIN64
BOOL UnsetTrayHook64 ()
#else
BOOL  UnsetTrayHook ()
#endif
{
    if( UnhookWindowsHookEx( g_hHook ) )
    {
        g_hHook = NULL;
        g_hTargetWin = NULL;
        return TRUE;
    }

    return FALSE;
}